#!/usr/bin/env python3
"""
DeepSpeed MOE + Qwen3-Coder API演示脚本
展示如何使用API进行各种编程任务
"""

import asyncio
import json
import time
from typing import List, Dict, Any
import aiohttp

class Qwen3CoderAPIDemo:
    """Qwen3-Coder API演示类"""
    
    def __init__(self, base_url: str = "http://localhost:8000", api_key: str = "roo-code-key-2024"):
        self.base_url = base_url
        self.api_key = api_key
        self.model = "qwen3-coder-480b-a35b-instruct"
    
    async def make_request(self, session: aiohttp.ClientSession, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
        """发起API请求"""
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        url = f"{self.base_url}/v1{endpoint}"
        
        async with session.post(url, json=data, headers=headers) as response:
            if response.status == 200:
                return await response.json()
            else:
                error_text = await response.text()
                raise Exception(f"API请求失败: {response.status} - {error_text}")
    
    async def health_check(self):
        """健康检查"""
        print("🔍 检查API服务状态...")
        
        async with aiohttp.ClientSession() as session:
            async with session.get(f"{self.base_url}/health") as response:
                if response.status == 200:
                    health_info = await response.json()
                    print(f"✅ 服务状态: {health_info['status']}")
                    print(f"🤖 模型状态: {'已加载' if health_info['model_loaded'] else '未加载'}")
                    return health_info
                else:
                    print("❌ API服务不可用")
                    return None
    
    async def demonstrate_code_generation(self):
        """代码生成演示"""
        print("\n" + "="*50)
        print("🎨 代码生成演示")
        print("="*50)
        
        tasks = [
            {
                "name": "排序算法",
                "prompt": "Write an efficient merge sort algorithm in Python with proper error handling and documentation.",
                "max_tokens": 400
            },
            {
                "name": "数据结构",
                "prompt": "Create a Python class for a priority queue with insert, delete, and peek operations.",
                "max_tokens": 350
            },
            {
                "name": "装饰器应用",
                "prompt": "Create a Python decorator that measures and logs the execution time of functions.",
                "max_tokens": 300
            }
        ]
        
        async with aiohttp.ClientSession() as session:
            for task in tasks:
                print(f"\n📝 生成: {task['name']}")
                print("-" * 30)
                
                data = {
                    "model": self.model,
                    "messages": [{"role": "user", "content": task['prompt']}],
                    "max_tokens": task['max_tokens'],
                    "temperature": 0.1
                }
                
                try:
                    start_time = time.time()
                    result = await self.make_request(session, "/chat/completions", data)
                    end_time = time.time()
                    
                    content = result['choices'][0]['message']['content']
                    usage = result['usage']
                    
                    print(f"⏱️  响应时间: {usage['latency_seconds']:.2f}秒")
                    print(f"📊 Token使用: {usage['total_tokens']}个")
                    print(f"🚀 速度: {usage['tokens_per_second']:.1f} tokens/s")
                    print(f"\n💻 生成代码:\n{content}")
                    
                except Exception as e:
                    print(f"❌ 错误: {e}")
    
    async def demonstrate_code_debugging(self):
        """代码调试演示"""
        print("\n" + "="*50)
        print("🐛 代码调试演示")
        print("="*50)
        
        buggy_code = '''
def factorial(n):
    if n == 0:
        return 1
    else:
        return n * factorial(n-1)

# 测试
for i in range(-1, 5):
    print(f"factorial({i}) = {factorial(i)}")
'''
        
        prompt = f"Please debug this code and explain the issues:\n```python\n{buggy_code}\n```"
        
        print("🔍 分析代码...")
        async with aiohttp.ClientSession() as session:
            data = {
                "model": self.model,
                "system": "You are an expert debugger. Identify issues and provide fixes.",
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": 500,
                "temperature": 0.2
            }
            
            try:
                result = await self.make_request(session, "/chat/completions", data)
                analysis = result['choices'][0]['message']['content']
                usage = result['usage']
                
                print(f"\n🔧 调试分析:\n{analysis}")
                print(f"\n📊 使用了 {usage['total_tokens']} 个tokens")
                
            except Exception as e:
                print(f"❌ 错误: {e}")
    
    async def demonstrate_code_review(self):
        """代码审查演示"""
        print("\n" + "="*50)
        print("🔍 代码审查演示")
        print("="*50)
        
        sample_code = '''
import requests
import json

def get_user_data(user_id):
    url = f"https://api.example.com/users/{user_id}"
    response = requests.get(url)
    if response.status_code == 200:
        return response.json()
    else:
        return None

def process_data(users):
    for user in users:
        data = get_user_data(user['id'])
        if data:
            print(f"Processing {data['name']}")
'''
        
        prompt = f"Please review this code for best practices, security, and performance:\n```python\n{sample_code}\n```"
        
        print("🔍 审查代码质量...")
        async with aiohttp.ClientSession() as session:
            data = {
                "model": self.model,
                "system": "You are a senior software engineer conducting a code review. Focus on security, performance, maintainability, and best practices.",
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": 600,
                "temperature": 0.3
            }
            
            try:
                result = await self.make_request(session, "/chat/completions", data)
                review = result['choices'][0]['message']['content']
                usage = result['usage']
                
                print(f"\n📋 代码审查报告:\n{review}")
                print(f"\n📊 使用了 {usage['total_tokens']} 个tokens")
                
            except Exception as e:
                print(f"❌ 错误: {e}")
    
    async def demonstrate_algorithm_explanation(self):
        """算法解释演示"""
        print("\n" + "="*50)
        print("🧠 算法解释演示")
        print("="*50)
        
        algorithms = [
            {
                "name": "快速排序",
                "code": '''
def quicksort(arr):
    if len(arr) <= 1:
        return arr
    pivot = arr[len(arr) // 2]
    left = [x for x in arr if x < pivot]
    middle = [x for x in arr if x == pivot]
    right = [x for x in arr if x > pivot]
    return quicksort(left) + middle + quicksort(right)
''',
                "description": "Explain the quicksort algorithm, its time complexity, and when to use it."
            },
            {
                "name": "二叉搜索树",
                "code": '''
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class BST:
    def search(self, root, target):
        if not root or root.val == target:
            return root
        if target < root.val:
            return self.search(root.left, target)
        return self.search(root.right, target)
''',
                "description": "Explain how binary search trees work and analyze their operations."
            }
        ]
        
        async with aiohttp.ClientSession() as session:
            for algo in algorithms:
                print(f"\n📚 分析: {algo['name']}")
                print("-" * 30)
                
                prompt = f"{algo['description']}\n```python\n{algo['code']}\n```"
                
                data = {
                    "model": self.model,
                    "system": "You are a computer science professor explaining algorithms clearly and comprehensively.",
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": 500,
                    "temperature": 0.4
                }
                
                try:
                    result = await self.make_request(session, "/chat/completions", data)
                    explanation = result['choices'][0]['message']['content']
                    usage = result['usage']
                    
                    print(f"📖 算法解释:\n{explanation}")
                    print(f"\n⏱️  响应时间: {usage['latency_seconds']:.2f}秒")
                    
                except Exception as e:
                    print(f"❌ 错误: {e}")
    
    async def demonstrate_streaming(self):
        """流式输出演示"""
        print("\n" + "="*50)
        print("🌊 流式输出演示")
        print("="*50)
        
        prompt = "Explain the differences between REST and GraphQL APIs, and when to use each."
        
        print(f"📝 询问: {prompt}")
        print("💬 回答中...")
        
        async with aiohttp.ClientSession() as session:
            data = {
                "model": self.model,
                "messages": [{"role": "user", "content": prompt}],
                "stream": True,
                "max_tokens": 800,
                "temperature": 0.3
            }
            
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }
            
            try:
                async with session.post(
                    f"{self.base_url}/v1/chat/completions", 
                    json=data, 
                    headers=headers
                ) as response:
                    if response.status == 200:
                        print("✅ 流式响应开始:")
                        print("-" * 40)
                        
                        full_response = ""
                        async for line in response.content:
                            line = line.decode('utf-8').strip()
                            if line.startswith('data: '):
                                data_str = line[6:]  # Remove 'data: ' prefix
                                if data_str == '[DONE]':
                                    break
                                
                                try:
                                    chunk_data = json.loads(data_str)
                                    if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
                                        if 'delta' in chunk_data['choices'][0] and 'content' in chunk_data['choices'][0]['delta']:
                                            content = chunk_data['choices'][0]['delta']['content']
                                            print(content, end='', flush=True)
                                            full_response += content
                                except json.JSONDecodeError:
                                    continue
                        
                        print("\n" + "-" * 40)
                        print(f"✅ 流式响应完成，总长度: {len(full_response)} 字符")
                        
                    else:
                        print(f"❌ 流式请求失败: {response.status}")
                        
            except Exception as e:
                print(f"❌ 错误: {e}")
    
    async def demonstrate_concurrent_requests(self):
        """并发请求演示"""
        print("\n" + "="*50)
        print("🚀 并发请求演示")
        print("="*50)
        
        tasks = [
            "Write a Python decorator for caching function results.",
            "Explain the SOLID principles in object-oriented design.",
            "Create a function to check if a string is a palindrome.",
            "Describe the differences between processes and threads.",
            "Write a SQL query to find duplicate records."
        ]
        
        print(f"📝 同时发送 {len(tasks)} 个请求...")
        
        async def single_request(task: str, index: int):
            async with aiohttp.ClientSession() as session:
                data = {
                    "model": self.model,
                    "messages": [{"role": "user", "content": task}],
                    "max_tokens": 200,
                    "temperature": 0.2
                }
                
                start_time = time.time()
                try:
                    result = await self.make_request(session, "/chat/completions", data)
                    end_time = time.time()
                    
                    response_time = end_time - start_time
                    tokens = result['usage']['total_tokens']
                    
                    return {
                        "index": index,
                        "task": task[:30] + "...",
                        "success": True,
                        "time": response_time,
                        "tokens": tokens
                    }
                except Exception as e:
                    return {
                        "index": index,
                        "task": task[:30] + "...",
                        "success": False,
                        "error": str(e)
                    }
        
        # 发送并发请求
        start_time = time.time()
        results = await asyncio.gather(*[
            single_request(task, i) for i, task in enumerate(tasks)
        ])
        total_time = time.time() - start_time
        
        # 显示结果
        print("\n📊 并发请求结果:")
        print("-" * 60)
        
        for result in results:
            if result["success"]:
                print(f"✅ 任务 {result['index']+1}: {result['task']} ({result['time']:.2f}s, {result['tokens']} tokens)")
            else:
                print(f"❌ 任务 {result['index']+1}: {result['task']} - {result['error']}")
        
        print("-" * 60)
        print(f"⏱️  总耗时: {total_time:.2f}秒")
        print(f"📈 平均每个请求: {total_time/len(tasks):.2f}秒")
        successful = sum(1 for r in results if r["success"])
        print(f"✅ 成功率: {successful}/{len(tasks)} ({successful/len(tasks)*100:.1f}%)")
    
    async def run_demo(self):
        """运行完整演示"""
        print("🎯 DeepSpeed MOE + Qwen3-Coder API演示")
        print("=" * 50)
        
        # 检查服务状态
        health = await self.health_check()
        if not health:
            print("❌ API服务不可用，请确保服务正在运行")
            return
        
        # 运行各项演示
        try:
            await self.demonstrate_code_generation()
            await self.demonstrate_code_debugging()
            await self.demonstrate_code_review()
            await self.demonstrate_algorithm_explanation()
            await self.demonstrate_streaming()
            await self.demonstrate_concurrent_requests()
            
            print("\n" + "="*50)
            print("🎉 演示完成！")
            print("="*50)
            print("\n💡 使用提示:")
            print("• API地址: http://localhost:8000")
            print("• 文档地址: http://localhost:8000/docs")
            print("• 健康检查: http://localhost:8000/health")
            print("• 统计信息: curl -H 'Authorization: Bearer roo-code-key-2024' http://localhost:8000/stats")
            
        except Exception as e:
            print(f"❌ 演示过程中发生错误: {e}")

async def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="DeepSpeed MOE + Qwen3-Coder API演示")
    parser.add_argument("--url", default="http://localhost:8000", help="API基础地址")
    parser.add_argument("--key", default="roo-code-key-2024", help="API密钥")
    
    args = parser.parse_args()
    
    demo = Qwen3CoderAPIDemo(args.url, args.key)
    await demo.run_demo()

if __name__ == "__main__":
    # 运行演示
    asyncio.run(main())